__author__ = 'shawhen'

import xml.etree.ElementTree as ET

ET.register_namespace("", "http://www.innotek.de/VirtualBox-settings")


class VboxConfig():
    def __init__(self, conf_file_path, auto_save=True):
        self._conf_file_path = conf_file_path
        self._auto_save = auto_save

        self._vbox_tree = ET.parse(conf_file_path)
        self._vbox_root = self._vbox_tree.getroot()

    def save(self):
        self._vbox_tree.write(self._conf_file_path, xml_declaration=True, default_namespace="")

    def set_machine_uuid(self, uuid):
        machine_ele = self._vbox_root.find(r"{http://www.innotek.de/VirtualBox-settings}Machine")
        if machine_ele is None:
            raise ValueError("could not found Machine Node")
        machine_ele.set("uuid", "{"+uuid+"}")
        if self._auto_save is True:
            self.save()

    def set_machine_name(self, name):
        machine_ele = self._vbox_root.find(r"{http://www.innotek.de/VirtualBox-settings}Machine")
        if machine_ele is None:
            raise ValueError("could not found Machine Node")
        machine_ele.set("name", name)
        if self._auto_save is True:
            self.save()

    def set_port_map(self, rule_name, proto="tcp", srcip="0.0.0.0", srcport=1,
                     dstip="0.0.0.0", dstport=1,
                     adapter_slot=0, mac_addr=None):
        if mac_addr is not None and adapter_slot != -1:
            raise ValueError("MacAddress is not None, so adapter slot should be -1")

        machine_ele = self._vbox_root.find(r"{http://www.innotek.de/VirtualBox-settings}Machine")
        if machine_ele is None:
            raise ValueError("could not found Machine Node")
        hardware_ele = machine_ele.find(r"{http://www.innotek.de/VirtualBox-settings}Hardware")
        if hardware_ele is None:
            raise ValueError("could not found Hardware Node")
        network_ele = hardware_ele.find(r"{http://www.innotek.de/VirtualBox-settings}Network")
        if network_ele is None:
            raise ValueError("could not found Network Node")
        adapter_eles = network_ele.findall("{http://www.innotek.de/VirtualBox-settings}Adapter")
        for adapter_ele in adapter_eles:
            if int(adapter_ele.get("slot")) == adapter_slot:
                break
            elif adapter_ele.get("MACAddress") == mac_addr:
                break
        else:
            raise ValueError("could not found Adapter(slot: {0}, MacAddress: {1}) Node".format(adapter_slot, mac_addr))
        nat_ele = adapter_ele.find("{http://www.innotek.de/VirtualBox-settings}NAT")
        if nat_ele is None:
            raise ValueError("could not found NAT Node")

        proto_str2value = {
            "udp": 0, "tcp": 1
        }
        forwarding_eles = nat_ele.findall("{http://www.innotek.de/VirtualBox-settings}Forwarding")
        for forwarding_ele in forwarding_eles:
            if forwarding_ele.get("name") == rule_name:  # 修改现有规则
                forwarding_ele.set("proto", str(proto_str2value[proto]))
                forwarding_ele.set("hostip", srcip)
                forwarding_ele.set("hostport", str(srcport))
                forwarding_ele.set("guestip", dstip)
                forwarding_ele.set("guestport", str(dstport))

                break
        else:  # 新增的一个规则
            forwarding_ele = ET.Element("Forwarding", {
                "name": rule_name,
                "proto": str(proto_str2value[proto]),
                "hostip": srcip,
                "hostport": str(srcport),
                "guestip": dstip,
                "guestport": str(dstport),
            })
            nat_ele.append(forwarding_ele)

        if self._auto_save is True:
            self.save()

    def set_mem_size(self, size_in_m):
        machine_ele = self._vbox_root.find(r"{http://www.innotek.de/VirtualBox-settings}Machine")
        if machine_ele is None:
            raise ValueError("could not found Machine Node")
        hardware_ele = machine_ele.find(r"{http://www.innotek.de/VirtualBox-settings}Hardware")
        if hardware_ele is None:
            raise ValueError("could not found Hardware Node")
        memory_ele = hardware_ele.find("{http://www.innotek.de/VirtualBox-settings}Memory")
        if memory_ele is None:
            raise ValueError("could not found Memory Node")
        memory_ele.set("RAMSize", str(size_in_m))

        if self._auto_save is True:
            self.save()